package cn.satoken.filter;

import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.Refill;
import org.springframework.stereotype.Component;

import javax.servlet.*;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.time.Duration;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;

/**
 * <h3>saToken</h3>
 * <p>ip限流防止脚本，频繁刷新验证码接口</p>
 *
 * @author gdw
 * @since 2025-06-06 10:44:01
 */
@Component
public class CaptchaRateLimiterFilter implements Filter {

    private final Cache<String, Bucket> bucketCache = CacheBuilder.newBuilder()
            .expireAfterAccess(10, TimeUnit.MINUTES)
            .maximumSize(10000)
            .build();

    private Bucket createNewBucket() {
        Refill refill = Refill.intervally(10, Duration.ofMinutes(1));
        Bandwidth limit = Bandwidth.classic(10, refill);
        return Bucket.builder().addLimit(limit).build();
    }

    private String getClientIP(HttpServletRequest request) {
        String xfHeader = request.getHeader("X-Forwarded-For");
        return (xfHeader != null && !xfHeader.isEmpty() && !"unknown".equalsIgnoreCase(xfHeader))
                ? xfHeader.split(",")[0].trim()
                : request.getRemoteAddr();
    }

    @Override
    public void doFilter(ServletRequest req, ServletResponse res, FilterChain chain)
            throws IOException, ServletException {

        HttpServletRequest request = (HttpServletRequest) req;
        HttpServletResponse response = (HttpServletResponse) res;

        if (request.getRequestURI().contains("/api/captcha")) {
            String ip = getClientIP(request);
            Bucket bucket = null;
            try {
                bucket = bucketCache.get(ip, this::createNewBucket);
            } catch (ExecutionException e) {
                throw new RuntimeException(e);
            }
            if (!bucket.tryConsume(1)) {
                response.setStatus(429);
                response.setContentType("application/json;charset=UTF-8");
                response.getWriter().write("{\"code\":429, \"msg\":\"请求过于频繁，请稍后再试。\"}");
                return;
            }
        }

        chain.doFilter(req, res);
    }
}
